import React, { useState, useEffect, useCallback } from 'react';
import { Modal, Tree, message, App as AntdApp, Space, Spin } from 'antd';
import { MenuOutlined } from '@ant-design/icons';
import type { Menu } from '@/api/system';
import { getMenuTree, getRoleMenus, assignRoleMenus } from '@/api/system';
import { useQuery } from '@tanstack/react-query';
import type { DataNode, TreeProps } from 'antd/es/tree';

const { DirectoryTree } = Tree;

interface MenuAssignDialogProps {
  visible: boolean;
  roleId: string | number;
  roleName: string;
  onClose: () => void;
  onSuccess?: () => void;
}

const MenuAssignDialog: React.FC<MenuAssignDialogProps> = ({
  visible,
  roleId,
  roleName,
  onClose,
  onSuccess,
}) => {
  const { modal } = AntdApp.useApp();
  const [selectedKeys, setSelectedKeys] = useState<React.Key[]>([]);
  const [loading, setLoading] = useState(false);
  const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);

  // 获取菜单树形数据
  const { data: menuTree, refetch: refetchMenuTree } = useQuery({
    queryKey: ['menu-tree'],
    queryFn: async () => {
      const response = await getMenuTree();
      return response;
    },
    enabled: visible,
  });

  // 获取角色已分配的菜单
  const { data: roleMenus, refetch: refetchRoleMenus } = useQuery({
    queryKey: ['role-menus', roleId],
    queryFn: async () => {
      const response = await getRoleMenus(roleId);
      return response;
    },
    enabled: visible,
  });

  // 转换菜单数据为树形结构
  const transformMenuToTree = (menus: Menu[]): DataNode[] => {
    const transformNode = (menu: Menu): DataNode => {
      return {
        title: menu.name,
        key: menu.id,
        children: menu.children ? menu.children.map(transformNode) : undefined,
      };
    };
    return menus.map(transformNode);
  };

  // 获取所有需要展开的节点key
  const getAllExpandedKeys = (nodes: DataNode[]): React.Key[] => {
    return nodes.reduce((keys: React.Key[], node) => {
      keys.push(node.key);
      if (node.children) {
        keys.push(...getAllExpandedKeys(node.children));
      }
      return keys;
    }, []);
  };

  // 初始化选中状态和展开状态
  useEffect(() => {
    if (menuTree && roleMenus) {
      setSelectedKeys(roleMenus.map(String));
      const treeData = transformMenuToTree(menuTree);
      setExpandedKeys(getAllExpandedKeys(treeData));
    }
  }, [menuTree, roleMenus]);

  // 处理树选择
  const handleTreeSelect: TreeProps['onSelect'] = (keys, info) => {
    setSelectedKeys(keys);
  };

  // 处理树勾选
  const handleTreeCheck: TreeProps['onCheck'] = (keys, info) => {
    setSelectedKeys(keys as React.Key[]);
  };

  // 处理分配菜单
  const handleAssignMenus = async () => {
    if (selectedKeys.length === 0) {
      message.warning('请至少选择一个菜单');
      return;
    }

    setLoading(true);
    try {
      await assignRoleMenus(roleId, selectedKeys);
      message.success('菜单分配成功');
      onSuccess?.();
      onClose();
    } catch (error) {
      message.error('菜单分配失败');
    } finally {
      setLoading(false);
    }
  };

  // 渲染树节点图标
  const renderTreeNode = (node: DataNode) => {
    return {
      ...node,
      icon: <MenuOutlined />,
    };
  };

  const treeData = menuTree ? transformMenuToTree(menuTree) : [];

  return (
    <Modal
      title={`分配菜单 - ${roleName}`}
      open={visible}
      onCancel={onClose}
      width={600}
      onOk={handleAssignMenus}
      confirmLoading={loading}
      destroyOnClose
    >
      <div style={{
        padding: '16px 0',
        height: '500px',
        overflow: 'auto',
        border: '1px solid #d9d9d9',
        borderRadius: '6px'
      }}>
        {menuTree ? (
          <Tree
            checkable
            selectable
            checkedKeys={selectedKeys}
            selectedKeys={selectedKeys}
            expandedKeys={expandedKeys}
            onExpand={setExpandedKeys}
            onSelect={handleTreeSelect}
            onCheck={handleTreeCheck}
            treeData={treeData.map(renderTreeNode)}
            height={450}
            showLine
            showIcon
            blockNode
          />
        ) : (
          <div style={{
            display: 'flex',
            justifyContent: 'center',
            alignItems: 'center',
            height: '100%'
          }}>
            <Spin size="large" tip="加载菜单数据..." />
          </div>
        )}
      </div>
    </Modal>
  );
};

export default MenuAssignDialog;